# Alexander F. Gazmararian
# afg2@princeton.edu

library(tidyverse)
library(here)
library(modelsummary)
library(kableExtra)
library(margins)
library(sandwich)
library(lmtest)
library(viridis)
library(gridExtra)

# load data
fair <- readRDS(here("data", "fair_conjoint.rds"))
fair <- subset(fair, !is.na(selected))

nat <- readRDS(here("data", "national_conjoint.rds"))

# set up
source(here("code", "fun", "table_spec.r"))
source(here("code", "fun", "fix_txt.r"))

# rename covariate names for fair for consistency

fair <- fair %>%
  rename(
    Community.Investment = `Community Investment`,
    Free.Retraining.Program = `Free Retraining Program`,
    Income.Support.During.Retraining = `Income Support During Retraining`,
    Worker.Retraining.Time = `Worker Retraining Time`,
    Benefit.Support = `Benefit Support`,
    Retrained.Worker.Salary = `Retrained Worker Salary`,
    Relocation.Support = `Relocation Support`
  )

# specify model
f <- selected ~ Community.Investment + Free.Retraining.Program + Income.Support.During.Retraining + 
  Worker.Retraining.Time + Benefit.Support + Retrained.Worker.Salary + Relocation.Support

# estimate differential effect of information treatment
info <- list()

# Fair sample
info[[1]] <- lm(update(f, ~ . * diversify), fair, weights = weights) # Fair sample
info[[2]] <- lm(update(f, ~ . * diversify), fair) # Fair sample + no sampling weights
info[[3]] <- lm(update(f, ~ . * diversify + fair + sex + gw_binary + college + ffemploy + age + income_med + white + partysummary), fair, weights = weights) # Fair sample + covariates

#National sample
info[[4]] <- lm(update(f, ~ . * diversify), nat, weights = weights)
info[[5]] <- lm(update(f, ~ . * diversify), nat) #National model + no sampling weights
info[[6]] <- lm(update(f, ~ . * diversify + partysummary + age + sex + college + black + hispanic + ff + fullemploy + income5), nat, weights = weights)

# Create SI table ----
info_exp <- here("output", "tables", "si_tab_9.1.txt")
modelsummary(
  info,
  fmt = 2,
  vcov = ~ ResponseId,
  gof_map = gm,
  escape = FALSE,
  stars = c("*"=.1,"**"=.05, "***"=.01),
  coef_map = c(
    "Community.InvestmentBroadband" = "Community Investment: Broadband",
    "Community.InvestmentHousing for new residents" = "Community Investment: Housing for new residents",
    "Free.Retraining.ProgramFor jobs in clean energy"="Free Retraining Program: For jobs in clean energy",
    "Free.Retraining.ProgramFor jobs in health care"="Free Retraining Program: For jobs in healthcare",
    "Income.Support.During.Retraining$400 per week for workers"="Income Support During Retraining: \\$400 per week for workers",
    "Income.Support.During.RetrainingNone"="Income Support During Retraining: None",
    "Worker.Retraining.Time2 years"="Worker Retraining Time: 2 years",
    "Worker.Retraining.Time3-6 months"="Worker Retraining Time: 3-6 months",
    "Benefit.SupportFunding for worker health care"="Benefit Support: Funding for worker health care",
    "Benefit.SupportFunding for worker pensions"="Benefit Support: Funding for worker pensions",
    "Retrained.Worker.Salary$100,000 per year" = "Retrained Worker Salary: \\$100,000 per year",
    "Retrained.Worker.Salary$50,000 per year" = "Retrained Worker Salary: \\$50,000 per year",
    "Relocation.SupportVoucher to cover moving expenses" = "Relocation Support: Voucher to cover moving expenses",
    "Community.InvestmentBroadband:diversify" = "Community Investment: Broadband x Information",
    "Community.InvestmentHousing for new residents:diversify" = "Community Investment: Housing for new residents x Information",
    "Free.Retraining.ProgramFor jobs in clean energy:diversify"="Free Retraining Program: For jobs in clean energy x Information",
    "Free.Retraining.ProgramFor jobs in health care:diversify"="Free Retraining Program: For jobs in healthcare x Information",
    "Income.Support.During.Retraining$400 per week for workers:diversify"="Income Support During Retraining: \\$400 per week for workers x Information",
    "Income.Support.During.RetrainingNone:diversify"="Income Support During Retraining: None x Information",
    "Worker.Retraining.Time2 years:diversify"="Worker Retraining Time: 2 years x Information",
    "Worker.Retraining.Time3-6 months:diversify"="Worker Retraining Time: 3-6 months x Information",
    "Benefit.SupportFunding for worker health care:diversify"="Benefit Support: Funding for worker health care x Informationn",
    "Benefit.SupportFunding for worker pensions:diversify"="Benefit Support: Funding for worker pensions x Information",
    "Retrained.Worker.Salary$100,000 per year:diversify" = "Retrained Worker Salary: \\$100,000 per year x Information",
    "Retrained.Worker.Salary$50,000 per year:diversify" = "Retrained Worker Salary: \\$50,000 per year x Information",
    "Relocation.SupportVoucher to cover moving expenses:diversify" = "Relocation Support: Voucher to cover moving expenses x Information"
  ),
  add_rows = data.frame(
    c("Sample Weights", "Covariates"),
    c("Yes", "No"),
    c("No", "No"),
    c("Yes", "Yes"),
    c("Yes", "No"),
    c("No", "No"),
    c("Yes", "Yes")
  ),
  output = "latex"
) %>%
  add_header_above(c(" " = 1, "Fossil Fuel" = 3, "National" = 3)) %>%
  cat(., file = info_exp)
fix_txt(info_exp)

# Create figure ----

# function to get the interaction b/c margins package won't cooperate
get_inter <- function(model, control, treat) {
  require(sandwich)
  
  coef.out <- coef(model)
  coef.out <- coef.out[grepl("Free.Retraining.ProgramFor jobs in clean energy", names(coef.out))]
  
  beta.noinfo <- coef.out[[1]]
  beta.info <- coef.out[[2]] + beta.noinfo
  
  vcov.out <- vcovCL(model, ~ ResponseId)
  
  vcov.noinfo <- vcov.out["Free.Retraining.ProgramFor jobs in clean energy", "Free.Retraining.ProgramFor jobs in clean energy"]
  
  vcov.inter <- vcov.out["Free.Retraining.ProgramFor jobs in clean energy:diversify", "Free.Retraining.ProgramFor jobs in clean energy:diversify"]
  
  vcov.cross <- vcov.out["Free.Retraining.ProgramFor jobs in clean energy", "Free.Retraining.ProgramFor jobs in clean energy:diversify"]
  
  se.noninfo <- sqrt(vcov.noinfo)
  se.info <- sqrt(vcov.noinfo + vcov.inter + 2 * vcov.cross)
  
  dof <- model$df.residual
  
  df.out <- data.frame(
    d = c(control, treat),
    est = c(beta.noinfo, beta.info),
    se = c(se.noninfo, se.info)
  )
  
  df.out$lb95 <- with(df.out, est + qt(.025, dof) * se) 
  df.out$ub95 <- with(df.out, est + qt(.975, dof) * se) 
  
  df.out$lb90 <- with(df.out, est + qt(.05, dof) * se) 
  df.out$ub90 <- with(df.out, est + qt(.95, dof) * se) 
  
  return(df.out)
}

ame.fair <- get_inter(info[[3]], "Control\nN=1524", "Information\nN=1450")
ame.fair$Sample <- "Fossil Fuels"
ame.nat <- get_inter(info[[6]], "Control\nN=5010", "Information\nN=5000")
ame.nat$Sample <- "National"
ame.out <- rbind(ame.fair, ame.nat)

d_amount <- .5
coeftest(info[[3]], vcov. = vcovCL(info[[3]], ~ ResponseId))
coeftest(info[[6]], vcov. = vcovCL(info[[6]], ~ ResponseId))

p.ame <-
  ame.out %>%
  ggplot(aes(x = d, y = est)) +
  geom_hline(yintercept = 0 , color = "grey") +
  geom_errorbar(aes(ymin= lb95, ymax = ub95), width = 0, position = position_dodge(d_amount)) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25, position = position_dodge(d_amount)) +
  geom_point(size = 3, position = position_dodge(d_amount), fill = "white", shape = 21) +
  theme_bw(base_size = 14) +
  facet_wrap(~Sample, scales = "free_x") +
  labs(x = "Treatment Condition", y = "Support for Policy with\nClean Energy Jobs Program") +
  geom_text(
    data = data.frame(d = c("Control\nN=1524", "Control\nN=5010"),
                      est = c(-.175,-.175),
                      Sample = c("Fossil Fuels","National"),
                      label = c("Difference: p<0.05", "Difference: p<0.12")),
    aes(label = label),
    hjust = -.25
  ) +
  scale_y_continuous(limits = c(-.2, .2)) +
  theme(
    panel.grid = element_blank(),
    legend.position = ""
  )
p.ame



p.ame <-
  ame.out %>%
  ggplot(aes(x = d, y = est)) +
  geom_hline(yintercept = 0 , color = "grey") +
  geom_errorbar(aes(ymin= lb95, ymax = ub95), width = 0, position = position_dodge(d_amount)) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25, position = position_dodge(d_amount)) +
  geom_point(size = 3, position = position_dodge(d_amount), fill = "white", shape = 21) +
  theme_bw(base_size = 14) +
  facet_wrap(~Sample, scales = "free_x") +
  labs(x = "Treatment Condition", y = "Support for Policy with\nClean Energy Jobs Program") +
  geom_text(
    data = data.frame(d = c("Control\nN=1524", "Control\nN=5010"),
                      est = c(-.175,-.175),
                      Sample = c("Fossil Fuels","National"),
                      label = c("Difference: p<0.05", "Difference: p<0.")),
    aes(label = label),
    hjust = -.25
  ) +
  scale_y_continuous(limits = c(-.2, .2)) +
  theme(
    panel.grid = element_blank(),
    legend.position = ""
  )


ame.weights <- get_inter(info[[1]], "Control\nN=1524", "Information\nN=1450")
ame.weights$Sample <- "Sample Weights"

ame.all <- get_inter(info[[3]], "Control\nN=1524", "Information\nN=1450")
ame.all$Sample <- "Controls + Sample Weights"

ame.out <- rbind(ame.weights, ame.all)

d_amount <- .5

p.ame <-
  ame.out %>%
  ggplot(aes(x = d, y = est, color = Sample, shape = Sample)) +
  geom_hline(yintercept = 0 , color = "grey") +
  geom_errorbar(aes(ymin= lb95, ymax = ub95), width = 0, position = position_dodge(d_amount)) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25, position = position_dodge(d_amount)) +
  geom_point(size = 3, position = position_dodge(d_amount), fill = "white") +
  theme_bw(base_size = 14) +
  scale_color_grey(na.translate = FALSE, start = 0, end = .5) +
  scale_shape_manual(values = c(21, 24)) +
  #facet_wrap(~Sample, scales = "free_x") +
  labs(x = "Treatment Condition", y = "Support for Policy with\nClean Energy Jobs Program",
       shape = "", color = "") +
  scale_y_continuous(limits = c(-.2, .2)) +
  theme(
    panel.grid = element_blank(),
    legend.position = c(.2, .9),
    legend.background = element_blank()
  )
p.ame
ggsave(
  p.ame,
  filename = here("output", "figures", "fig_2.png"),
  width = 5.5,
  height = 2.8,
  scale = 1.5,
  dpi = 300
)




# Heterogeneous effects of information ----

het <- list()


# Fair sample
het[[1]] <- lm(update(f, ~ . * diversify * ffemploy), fair, weights = weights)
het[[2]] <- lm(update(f, ~ . * diversify * partysummary), fair, weights = weights)
het[[3]] <- lm(update(f, ~ . * diversify * pass_attn), fair, weights = weights)
het[[4]] <- lm(update(f, ~ . * diversify * college), fair, weights = weights)

# Fossil fuel employment
het.ffemploy <- 
  margins_summary(
    het[[1]],
    variables = "Free.Retraining.Program",
    at = list(diversify = 0:1, ffemploy = 0:1),
    vcov = vcovCL(het[[1]], ~ ResponseId)
  )
het.ffemploy <- data.frame(het.ffemploy)

ame.ff <-
  het.ffemploy %>%
  filter(grepl("clean energy", factor)) %>%
  mutate(ffemploy = ifelse(ffemploy == 1, "Fossil Fuels", "Non-Fossil Fuels"),
         diversify = ifelse(diversify == 1, "Information", "Control"),
         lb90 = AME + SE * qnorm(.05),
         ub90 = AME + SE * qnorm(.95)) %>%
  ggplot(aes(x = diversify, color = ffemploy, y = AME)) +
  geom_hline(yintercept = 0, color = "grey") +
  labs(x = "Treatment Condition", y = "Support for Policy with\nClean Energy Jobs Program",
       color = "", shape = "", title = "Household Fossil Fuel Employment") +
  geom_errorbar(aes(ymin = lower, ymax = upper), width = 0, position = position_dodge(.5)) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25, position = position_dodge(.5)) +
  geom_point(size = 3, position = position_dodge(.5), shape = 21, fill = "white") +
  scale_color_viridis(discrete = TRUE) +
  theme_bw(base_size = 14) +
  theme(
    panel.grid = element_blank(),
    legend.position = "bottom",
    legend.box.background = element_rect(color = "black")
  )
ame.ff

# Party Identification
het.pid <-
  margins_summary(
    het[[2]],
    variables = "Free.Retraining.Program",
    at = list(diversify = 0:1, partysummary = c("Democrat", "Neither", "Republican")),
    vcov = vcovCL(het[[2]], ~ ResponseId)
  )

modelsummary(het[[2]], vcov = ~ ResponseId, stars = TRUE)

het.pid <- data.frame(het.pid)
ame.pid <-
  het.pid %>%
  filter(grepl("clean energy", factor)) %>%
  mutate(diversify = ifelse(diversify == 1, "Information", "Control"),
         lb90 = AME + SE * qnorm(.05),
         ub90 = AME + SE * qnorm(.95)) %>%
  ggplot(aes(x = diversify, color = partysummary, y = AME)) +
  geom_hline(yintercept = 0, color = "grey") +
  labs(x = "Treatment Condition", y = "Support for Policy with\nClean Energy Jobs Program",
       color = "", shape = "", title = "Partisan Identification") +
  geom_errorbar(aes(ymin = lower, ymax = upper), width = 0, position = position_dodge(.5)) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25, position = position_dodge(.5)) +
  geom_point(size = 3, position = position_dodge(.5), shape = 21, fill = "white") +
  scale_color_viridis(discrete = TRUE) +
  theme_bw(base_size = 14) +
  theme(
    panel.grid = element_blank(),
    legend.position = "bottom",
    legend.box.background = element_rect(color = "black")
  )
ame.ff

# By attention
het.attn <-
  margins_summary(
    het[[3]],
    variables = "Free.Retraining.Program",
    at = list(diversify = 0:1, pass_attn = 0:1),
    vcov = vcovCL(het[[3]], ~ ResponseId)
  )
het.attn <- data.frame(het.attn)
ame.attn <-
  het.attn %>%
  filter(grepl("clean energy", factor)) %>%
  mutate(pass_attn = ifelse(pass_attn == 1, "More Attentive", "Less Attentive"),
         diversify = ifelse(diversify == 1, "Information", "Control"),
         lb90 = AME + SE * qnorm(.05),
         ub90 = AME + SE * qnorm(.95)) %>%
  ggplot(aes(x = diversify, color = pass_attn, y = AME)) +
  geom_hline(yintercept = 0, color = "grey") +
  labs(x = "Treatment Condition", y = "Support for Policy with\nClean Energy Jobs Program",
       color = "", shape = "", title = "Attentiveness") +
  geom_errorbar(aes(ymin = lower, ymax = upper), width = 0, position = position_dodge(.5)) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25, position = position_dodge(.5)) +
  geom_point(size = 3, position = position_dodge(.5), shape = 21, fill = "white") +
  scale_color_viridis(discrete = TRUE) +
  theme_bw(base_size = 14) +
  theme(
    panel.grid = element_blank(),
    legend.position = "bottom",
    legend.box.background = element_rect(color = "black")
  )
ame.attn

# By education
het.edu <-
  margins_summary(
    het[[4]],
    variables = "Free.Retraining.Program",
    at = list(diversify = 0:1, college = 0:1),
    vcov = vcovCL(het[[4]], ~ ResponseId)
  )
het.edu <- data.frame(het.edu)
ame.edu <-
  het.edu %>%
  filter(grepl("clean energy", factor)) %>%
  mutate(college = ifelse(college == 1, "College Degree", "No College Degree"),
         diversify = ifelse(diversify == 1, "Information", "Control"),
         lb90 = AME + SE * qnorm(.05),
         ub90 = AME + SE * qnorm(.95)) %>%
  ggplot(aes(x = diversify, color = college, y = AME)) +
  geom_hline(yintercept = 0, color = "grey") +
  labs(x = "Treatment Condition", y = "Support for Policy with\nClean Energy Jobs Program",
       color = "", shape = "", title = "Education") +
  geom_errorbar(aes(ymin = lower, ymax = upper), width = 0, position = position_dodge(.5)) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25, position = position_dodge(.5)) +
  geom_point(size = 3, position = position_dodge(.5), shape = 21, fill = "white") +
  scale_color_viridis(discrete = TRUE) +
  theme_bw(base_size = 14) +
  theme(
    panel.grid = element_blank(),
    legend.position = "bottom",
    legend.box.background = element_rect(color = "black")
  )
ame.edu

# Create plot for output
ame.het <- grid.arrange(ame.ff, ame.pid, ame.attn, ame.edu)

ggsave(
  ame.het,
  filename = here("output", "figures", "si_fig_9.1.png"),
  width = 6.5,
  height = 5,
  scale = 1.5,
  dpi = 300
)
